import numpy as np
from matplotlib import pyplot as plt
from sklearn.ensemble import IsolationForest

from Visualization.ReportArtist import create_report


class AnomalyHunter:
    def __init__(self, n_trees=100, contamination=0.1):
        """
        参数说明：
        n_trees : 森林中树的数量 -> 猎犬数量
        contamination : 数据中异常数据的比例
        """
        self.pack = IsolationForest(n_estimators=n_trees, contamination=contamination)

    def train(self, healthy_features):
        """
        训练巡警识别正常气味
        features : 健康样本特征矩阵
        """
        self.pack.fit(healthy_features)

    def detect(self, test_features):
        """
        执行异常嗅探
        返回：异常概率（0正常，1异常）
        """
        pred = self.pack.predict(test_features)
        return (pred == -1).astype(int)  # -1表示异常，转换为1

    def plot_heatmap(self, features):
        """生成损伤热点地图"""
        scores = -self.pack.score_samples(features)  # 获取每个样本的异常评分
        try:
            plt.figure(figsize=(8, 6))
            # 这里假设 scores 数组大小为 400，将其重塑为 (20, 20)
            plt.imshow(scores.reshape(20, 20), cmap='hot', interpolation='nearest')
            plt.colorbar(label='异常评分')
            plt.title("结构损伤热力图 - 越红越可疑")
            plt.xlabel("X 轴节点")
            plt.ylabel("Y 轴节点")
            plt.savefig('damage_heatmap.png')
            plt.close()
        except Exception as e:
            print(f"生成热力图失败: {e}")
        # 计算损伤概率（示例：取异常评分）
        damage_prob = (scores - np.min(scores)) / (np.max(scores) - np.min(scores))
        fig, diagnosis = create_report(
            damage_prob.reshape(20, 20).flatten(),  # 修改为 (20, 20)
            conclusion="结构损伤检测完成，请查看热力图详情"
        )
